/**
 * @file main.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "data_utils.h"
#include <cstring>

constexpr uint32_t SIG_HALF = 0;    // float16
constexpr uint32_t SIG_FLOAT = 1;   // float32
constexpr uint32_t SIG_INT16 = 2;   // int16_t
constexpr uint32_t SIG_INT32 = 3;   // int32_t
constexpr uint32_t INPUT_TYPE_SIZE[] = {2, 4, 2, 4};

constexpr uint32_t USE_CORE_NUM = 8;
constexpr uint32_t ALIGN_ELEM = 8;
constexpr uint32_t ALIGN_BASE = USE_CORE_NUM * ALIGN_ELEM;  // 对齐基准：64

#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void linspaced_custom_do(uint32_t blockDim, 
                                void *stream, 
                                uint8_t *start, 
                                uint8_t *end, 
                                const uint32_t n, 
                                uint8_t *z);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void linspaced_custom(GM_ADDR start, 
                                                       GM_ADDR end, 
                                                       const uint32_t n, 
                                                       GM_ADDR z);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;          
    size_t nByteSize = sizeof(uint32_t);                        
    size_t inputByteSize = INPUT_TYPE_SIZE[SIG_FLOAT];  

    uint32_t *ns = (uint32_t *)malloc(nByteSize);
    
    ReadFile("./input/ns.bin", nByteSize, ns, nByteSize); 
    
    const uint32_t n = ns[0];                

    const size_t totalAlignedSize = ((n + ALIGN_BASE - 1) / ALIGN_BASE) * ALIGN_BASE;  // (632→640)
    const size_t outputByteSizeAligned = totalAlignedSize * sizeof(float);             // 对齐后的内存大小（640×4=2560B）
    const size_t outputByteSizeOrig = n * sizeof(float);                               // 原始有效数据大小（632×4=2528B）

    free(ns);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *starts = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *ends = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSizeAligned);  
    
    ReadFile("./input/starts.bin", inputByteSize, starts, inputByteSize);
    ReadFile("./input/ends.bin", inputByteSize, ends, inputByteSize);
    
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    
    ICPU_RUN_KF(linspaced_custom, blockDim, starts, ends, n, z); 
    
    WriteFile("./output/output_z.bin", z, outputByteSizeOrig);
    
    AscendC::GmFree((void *)starts);
    AscendC::GmFree((void *)ends);
    AscendC::GmFree((void *)z);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
    
    uint8_t *startHost, *endHost, *zHost;       // 主机端内存（对齐后）
    uint8_t *startDevice, *endDevice, *zDevice; // 设备端内存（对齐后）
    
    CHECK_ACL(aclrtMallocHost((void **)(&startHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&endHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)&zHost, outputByteSizeAligned));   
    CHECK_ACL(aclrtMalloc((void **)&startDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&endDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSizeAligned, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/starts.bin", inputByteSize, startHost, inputByteSize);
    ReadFile("./input/ends.bin", inputByteSize, endHost, inputByteSize);
    
    CHECK_ACL(aclrtMemcpy(startDevice, inputByteSize, startHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(endDevice, inputByteSize, endHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    
    linspaced_custom_do(blockDim, stream, startDevice, endDevice, n, zDevice);
    CHECK_ACL(aclrtSynchronizeStream(stream));
    
    CHECK_ACL(aclrtMemcpy(zHost, outputByteSizeAligned, zDevice, outputByteSizeAligned, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, outputByteSizeOrig);

    CHECK_ACL(aclrtFree(startDevice));
    CHECK_ACL(aclrtFree(endDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(startHost));
    CHECK_ACL(aclrtFreeHost(endHost));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}